!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!! Finite state Markov-chain approximations to highly persistent processes
!! Karen A. Kopecky, Richard M.H. Suen
!! Review of Economic Dynamics
!! 13 (2010) 701-714
!! the error term has mean zero and standard error sig.
!! if the error term has mean mu, then shift the grid by +mu/(1-rho). everything else same.
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
MODULE ROUWENHORST
USE nr, ONLY: bico
IMPLICIT NONE
   INTEGER :: ii

   INTERFACE ar1
      MODULE PROCEDURE ar1_stat,ar1_plain
   END INTERFACE
CONTAINS

SUBROUTINE ar1_plain(rho,sig,grids,transmat)
!! NOTE: SIG IS THE CONDITIONAL SD OF THE RV
IMPLICIT NONE
   REAL(8),INTENT(IN) :: rho,sig
   REAL(8),INTENT(OUT) :: grids(:),transmat(:,:)

   INTEGER :: ns,i 
   REAL(8) :: psi,p,stp
   REAL(8),ALLOCATABLE :: temp(:,:),temp1(:,:)

   ns = SIZE(grids)

   ! rouwenhorst param and bounds
   p = (1d0+rho)/2d0
   psi = SQRT( (ns-1)/(1d0-rho**2) )*sig

   ! grid points
   stp = 2d0*psi/(ns-1)
   DO ii = 1,ns
      grids(ii) = -psi + stp*(ii-1)
   ENDDO

   ! transition matrix
   ALLOCATE(temp(2,2))
   temp(1,:) = (/p,1d0-p/)
   temp(2,:) = (/1d0-p,p/)

   DO ii = 3,ns
      ALLOCATE(temp1(ii,ii))
      temp1 = 0d0
      temp1(:ii-1,:ii-1) = p*temp
      temp1(:ii-1,2:) = temp1(:ii-1,2:)+(1d0-p)*temp
      temp1(2:,:ii-1) = temp1(2:,:ii-1)+(1d0-p)*temp
      temp1(2:,2:) = temp1(2:,2:)+p*temp

      temp1(2:ii-1,:) = temp1(2:ii-1,:)/2d0

      DEALLOCATE(temp); ALLOCATE(temp(ii,ii))
      temp = temp1
      DEALLOCATE(temp1)
   ENDDO
   transmat = temp
   DEALLOCATE(temp)
   RETURN
END SUBROUTINE

SUBROUTINE ar1_stat(rho,sig,grids,transmat,statdist)
!! NOTE: SIG IS THE CONDITIONAL SD OF THE RV
IMPLICIT NONE
   REAL(8),INTENT(IN) :: rho,sig
   REAL(8),INTENT(OUT) :: grids(:),transmat(:,:),statdist(:)

   INTEGER :: ns,i 
   REAL(8) :: psi,p,s,stp
   REAL(8),ALLOCATABLE :: temp(:,:),temp1(:,:)

   ns = SIZE(grids)

   ! rouwenhorst param and bounds
   p = (1d0+rho)/2d0
   psi = SQRT( (ns-1)/(1d0-rho**2) )*sig

   ! grid points
   stp = 2d0*psi/(ns-1)
   DO ii = 1,ns
      grids(ii) = -psi + stp*(ii-1)
   ENDDO

   ! transition matrix
   ALLOCATE(temp(2,2))
   temp(1,:) = (/p,1d0-p/)
   temp(2,:) = (/1d0-p,p/)

   DO ii = 3,ns
      ALLOCATE(temp1(ii,ii))
      temp1 = 0d0
      temp1(:ii-1,:ii-1) = p*temp
      temp1(:ii-1,2:) = temp1(:ii-1,2:)+(1d0-p)*temp
      temp1(2:,:ii-1) = temp1(2:,:ii-1)+(1d0-p)*temp
      temp1(2:,2:) = temp1(2:,2:)+p*temp

      temp1(2:ii-1,:) = temp1(2:ii-1,:)/2d0

      DEALLOCATE(temp); ALLOCATE(temp(ii,ii))
      temp = temp1
      DEALLOCATE(temp1)
   ENDDO
   transmat = temp
   DEALLOCATE(temp)

   s = (1d0-p) / (2d0*(1d0-p))
   DO ii = 1,ns
      statdist(ii) = bico(ns-1,ii-1)*s**(ns-ii)*(1d0-s)**(ii-1)
   ENDDO

   RETURN
END SUBROUTINE

END MODULE
